package com.cddx.logistics.controller;

import com.cddx.logistics.service.OllamaService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Map;

@RestController
@RequestMapping("/oll/ai")
public class AIController {

    @Autowired
    private OllamaService ollamaService;

    private boolean isDatabaseQuery(String question) {
        // 简单判断是否为数据库查询
        return question.contains("计划") ;
                //|| question.contains("库存") || question.contains("用户信息");
    }

    /**
     * 流式传输接口
     * 处理用户提问并返回流式响应
     */
    @PostMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamQuery(@RequestBody Map<String, String> payload) {

        String question = payload.get("question");
        if (isDatabaseQuery(question)) {
            return Flux.just(ollamaService.handleDatabaseQuery(question));
        }
        System.out.println(question);
        return ollamaService.streamQuery(question);
    }

    /**
     * 停止回答接口
     * 用于中断当前正在进行的流式响应
     */
    @PostMapping("/stop")
    public ResponseEntity<Void> stopAnswering() {
        // 这里需要实现中断逻辑，具体实现取决于你的服务架构
        // 如果是基于请求级别的中断，可以通过存储和管理每个请求的AbortController来实现
        return ResponseEntity.ok().build();
    }

    /**
     * 清空对话接口
     * 用于重置对话历史
     */
    @PostMapping("/clear")
    public ResponseEntity<Void> clearConversation() {
        // 如果需要在服务器端维护对话历史，可以在这里实现清除逻辑
        // 对于纯前端实现的对话历史，这个接口可以只是一个占位符
        return ResponseEntity.ok().build();
    }

    /**
     * 同步查询接口
     * 作为流式传输的备选方案，返回完整响应
     */
    @PostMapping("/query")
    public ResponseEntity<String> queryModel(@RequestBody Map<String, String> payload) {
        String question = payload.get("question");
        String response = ollamaService.queryModel(question);
        return ResponseEntity.ok(response);
    }
}
